import shutil
import os

import pytest
from fastapi.testclient import TestClient

from main import create_app

from search.kg import ChronRAG
from load.checkpoint import Checkpoint
from llm import LLM

@pytest.fixture(name="client")
def client_fixture():
    app = create_app()

    folder_path = "/tmp/randm129temp_folder_for_cktest"
    os.makedirs(folder_path, exist_ok=True)

    with TestClient(app=app) as client:
        yield client

    if os.path.exists(folder_path):
        shutil.rmtree(folder_path)        

@pytest.fixture(name="dst")
def dst_create():
    return ChronRAG()

@pytest.fixture(name="cp")
def cp_create():
    # 创建临时文件夹路径
    folder_path = "/tmp/randm129temp_folder_for_cp"
    os.makedirs(folder_path, exist_ok=True)

    yield Checkpoint(folder_path)

    if os.path.exists(folder_path):
        shutil.rmtree(folder_path)

@pytest.fixture(name="llm")
def llm_create():
    return LLM()
